# coding:utf-8
import sys
import math

def load_data(data_path):


    """
    
    :param data_path: 训练数据文件的路径
    :return: 
        all_wordcount_dict = {word:count} 所有文章单词未去重，每个单词词频
        wordcount_for_class_dict={classid:{word:wordcount}} 某一类别中的文章中，每个单词的词频
        wordcount_distict_dict= {word:1} 所有文章中去重，单词的总数，每个单词是 1
        class_all_word_count = {classid:wordcount} 某一类文章中，单词的总数
        classcount_dict={classid:count} 分类名称:这个类别有多少篇文章
    """
    all_wordcount_dict = {}
    wordcount_for_class_dict = {}
    wordcount_distict_dict= {}
    class_all_word_count = {}
    classcount_dict = {}
    with open(data_path) as fd:
        for line in fd.readlines():
            line = line.strip()
            # 去掉 "#" 及后面所有的字符
            ops = line.find("#")
            if ops > 0:
                line = line[0:ops].strip()
            else:
                continue

            word_list = line.split(" ")
            # 第一个字符是，类id
            classid = word_list[0]
            word_list = word_list[1:]

            if classid not in classcount_dict:
                classcount_dict[classid] = 1
            else:
                classcount_dict[classid] += 1

            if classid not in class_all_word_count:
                class_all_word_count[classid] = 0
            for word in word_list:
                if len(word) < 1:
                    print("file format error")
                    continue
                class_all_word_count[classid] += 1
                if word not in wordcount_distict_dict:
                    wordcount_distict_dict[word] = 1
                if word not in all_wordcount_dict:
                    all_wordcount_dict[word] = 1
                else:
                    all_wordcount_dict[word] += 1

                if classid not in wordcount_for_class_dict:
                    wordcount_for_class_dict[classid] = {}
                else:
                    if word not in wordcount_for_class_dict[classid]:
                        wordcount_for_class_dict[classid][word] = 1
                    else:
                        wordcount_for_class_dict[classid][word] += 1


    return all_wordcount_dict, wordcount_for_class_dict, wordcount_distict_dict, class_all_word_count, classcount_dict

def ComputeModel(all_wordcount_dict, wordcount_for_class_dict, wordcount_distict_dict,
                 class_all_word_count,classcount_dict,defaultFreq=0.1):

    """
    
    :param all_wordcount_dict: 
    :param wordcount_for_class_dict: 
    :param wordcount_distict_dict: 
    :param class_all_word_count: 
    :param classcount_dict: 
    :param defaultFreq: 
    :return: 
        word_prob_for_class_dict = {classid:{word:p(word|classid)}} 某一类文章中，每个单词的概率
        class_default_prob_dict = {classid:p(default_prob)} 每个类别中，如果没有出现过这个单词，人为的给一个概率
        class_prob_dict = {classid:p(classid)} 先验概率，每一类文章的概率
    
    p(word|classid) = (classid 类文章中，word出现的次数 + defaultFreq)/((所有文章对单词去重后单词的总数)*defaultFreq+classid 类文文章中单词的总数）
    """
    word_prob_for_class_dict = {}
    # 计算某一类文章的总数
    class_sum = sum(classcount_dict.values())
    # 计算去重后，单词的总数
    all_word_sum = len(wordcount_distict_dict)
    class_default_prob_dict = {}
    class_prob_dict = {}
    for classid in wordcount_for_class_dict:

        class_prob_dict[classid] = float(classcount_dict[classid]) / float(class_sum)
        word_prob_for_class_dict[classid] = {}
        class_all_word_sum = class_all_word_count[classid]
        sum_count = float(all_word_sum * defaultFreq + class_all_word_sum)
        class_default_prob_dict[classid] = float(defaultFreq) / sum_count

        for word in wordcount_for_class_dict[classid]:
            class_word_count = float(wordcount_for_class_dict[classid][word] + defaultFreq)
            word_prob_for_class_dict[classid][word] = class_word_count / sum_count

    return word_prob_for_class_dict, class_default_prob_dict, class_prob_dict

def saveModel(word_prob_for_class_dict,class_default_prob_dict,class_prob_dict,wordcount_for_class_dict,model_file_path):
    """
    将模型保存在文件中
    :param word_prob_for_class_dict: 
    :param class_default_prob_dict: 
    :param class_prob_dict: 
    :param wordcount_for_class_dict: 
    :param model_file_path: 模型文件的路径
    :return: 
    """
    with open(model_file_path,"w",encoding="UTF-8") as fdModel:
        # 第一行内容：classid p(classid) class_default_prob_dict[classid]
        for classid in class_prob_dict:
            fdModel.write(classid + " ")
            fdModel.write(str(class_prob_dict[classid]) + " ")
            fdModel.write(str(class_default_prob_dict[classid]) + " ")

        fdModel.write("\n")

        # 每一个 classid 一行 word1 p(word1) word2 p(word2) ..... wordn p(wordn)
        for classid in wordcount_for_class_dict:
            for word in wordcount_for_class_dict[classid]:
                fdModel.write(word + " " +str(word_prob_for_class_dict[classid][word]))
                fdModel.write(" ")
            fdModel.write("\n")

def Predict(word_prob_for_class_dict,class_default_prob_dict,class_prob_dict,wordcount_distict_dict,test_file_path):
    """
    做预测，正常情况下预测是需要加载模型的，而我这里预测和计算模型在同一个文件中，为了性能就不加载模型了，如果模型过大，可能还是需要从文件中加载
    :param word_prob_for_class_dict: 
    :param class_default_prob_dict: 
    :param class_prob_dict: 
    :param wordcount_distict_dict: 
    :param test_file_path: 
    :return: 
        trueLabelList 存储真实标签 [classid1,classid2,....,classidn] 有很多重复的值
        predictLabelList  存储预测标签
    """
    i = 0
    classid_prob_dict = {}
    trueLabelList = []
    predictLabelList = []
    with open(test_file_path,"r",encoding="UTF-8") as fd_test:
        for line in fd_test.readlines():
            line = line.strip()
            ops = line.find("#")
            if ops > 0:
                line = line[:ops]
            else:
                continue
            word_list = line.split(" ")
            classid = word_list[0]
            trueLabelList.append(classid)
            word_list = word_list[1:]
            for classid in class_prob_dict:
                # 这里求 log 是想让最后的概率不至于太小，而造成精度越界，对预测结果没有影响，预测结果主要看那类概率大就判断成哪一类
                classid_prob_dict[classid] = math.log(class_prob_dict[classid])
            for word in word_list:
                if len(word) < 1:
                    continue
                if word not in wordcount_distict_dict:
                    continue
                for classid in class_prob_dict:
                    if word not in word_prob_for_class_dict[classid]:
                        classid_prob_dict[classid] += math.log(class_default_prob_dict[classid])
                    else:
                        classid_prob_dict[classid] += math.log(word_prob_for_class_dict[classid][word])
            max_prob = max(classid_prob_dict.values())
            for classid in classid_prob_dict:
                if classid_prob_dict[classid] == max_prob:
                    predictLabelList.append(classid)

    return trueLabelList, predictLabelList

def Evaluate(trueLabelList,predictLabelList):
    """
    计算准确度
    :param trueLabelList: 
    :param predictLabelList: 
    :return: 
    """
    accuracy_sum = 0
    for i in range(len(trueLabelList)):
        if trueLabelList[i] == predictLabelList[i]:
            accuracy_sum += 1
    print("this model accurary: ",float(accuracy_sum) / float(len(trueLabelList)))

def CalPreRec(trueLabelList,predictLabelList,classcount_dict):
    """
    计算召回率和精确率
    :param trueLabelList: 
    :param predictLabelList: 
    :param classcount_dict: 
    :return: 
    """
    tp = 0
    precision_sum = 0
    recall_sum = 0
    for classid in classcount_dict:
        for i in range(len(trueLabelList)):
            if trueLabelList[i] == classid and predictLabelList[i] == classid:
                tp += 1
            if predictLabelList[i] == classid:
                precision_sum += 1
            if trueLabelList[i] == classid:
                recall_sum += 1
        print("classid:{} precision is {}".format(classid,str(float(tp) / float(precision_sum))))
        print("classid:{} recall is {}".format(classid,str(float(tp) / float(recall_sum))))

def main():
    try:
        train_file_path = '/data/project/python/code/machine_learning/nb/nb_test/nb_data.train'
        test_file_path = '/data/project/python/code/machine_learning/nb/nb_test/nb_data.test'
        model_file_path = sys.argv[1]
        # print(sys.argv[0])
        if len(sys.argv) < 1:
            print("Use" + sys.argv[0] + "model_file_path")
        all_wordcount_dict, wordcount_for_class_dict, wordcount_distict_dict, class_all_word_count, classcount_dict = \
            load_data(train_file_path)
        word_prob_for_class_dict, class_default_prob_dict, class_prob_dict = \
            ComputeModel(all_wordcount_dict, wordcount_for_class_dict, wordcount_distict_dict, class_all_word_count, classcount_dict)
        saveModel(word_prob_for_class_dict, class_default_prob_dict, class_prob_dict, wordcount_for_class_dict,
                  model_file_path)
        trueLabelList,predictLabelList = Predict(word_prob_for_class_dict, class_default_prob_dict, class_prob_dict, wordcount_distict_dict,test_file_path)
        Evaluate(trueLabelList, predictLabelList)
        CalPreRec(trueLabelList, predictLabelList, classcount_dict)
    except IndexError:
        print("Use" + sys.argv[0] + "model_file_path")

if __name__ == '__main__':
    main()